from utility.algo2 import Algorithm
from utility.objects import *
import pandas as pd
from datetime import datetime, timedelta
from statsmodels.tsa.stattools import coint, adfuller

from pathlib import Path
import os

class PairTrading(Algorithm):

    init_ratio = 1
    steps = 1
    gap = 0.01
    init_eq = 10000
    order_size = 100

    symbol_1 = ''
    symbol_2 = ''

    def __init__(self, name:str, params:dict, begin :datetime, end:datetime):
        self.name = name
        self.begin = begin
        self.end = end
        for param in params:
            setattr(self, param, params[param])

        super().__init__('pair','symbol',params,begin,end)



    def run(self):
        msg = "running strategy" + self.name
        print(msg)
        data_1 = self.symbol_1.data.mkt_day
        data_2 = self.symbol_2.data.mkt_day

        ratio = pd.DataFrame(data_1.close / data_2.close)
        vol_window = self.vol_window
        mean_window = self.mean_window

        ratio['volatility'] = ratio['close'].rolling(window=vol_window).std() / np.sqrt(vol_window)
        # merge data and calculate ratio and vol
        ratio['smooth_vol'] = ratio['volatility'].rolling(window=mean_window).mean()

        model_window = self.model_window
        trade_window = self.trade_window

        # now we model the ratio
        def adf_pvalue(arr):
            """
            传入 1D ndarray，先去均值再做 ADF，返回 p‑value
            用 maxlag=0 避免窗口太短自选滞后失败。
            """
            series = arr - arr.mean()
            try:
                return round(adfuller(series, maxlag=0, regression='c', autolag=None)[1],4)
            except ValueError:  # 全常数或样本太少
                return np.nan
        ratio['adf_value'] = (ratio['close'].rolling(window = model_window).apply(adf_pvalue))

        ratio.dropna(inplace=True)
        print(ratio)
        # 固定p值启动策略
        # 0.2

        # init trading parameters
        self.mid_ratio = None
        self.current_step = 0
        self.step_gap = 0
        self.posn_1 = 0
        self.posn_2 = 0
        status = ''
        # 不能这么写，得写两个algo，组成一个portfolio，两个algo反着做，然后一个策略当成一个portfolio.

        for day, row in ratio.iterrows():
            print(row)
            next_day_begin = day + timedelta(days=1)
            next_day_end = next_day_begin + timedelta(hours=23, minutes=59, seconds=59)
            mkt_data_1 = self.symbol_1.data.mkt_data[next_day_begin:next_day_end]
            mkt_data_1 = mkt_data_1.add_suffix('_1')
            mkt_data_2 = self.symbol_2.data.mkt_data[next_day_begin:next_day_end]
            mkt_data_2 = mkt_data_2.add_suffix('_2')

            merged_data = pd.merge(mkt_data_1, mkt_data_2, on='datetime')
            merged_data['ratio'] = merged_data['close_1'] / merged_data['close_2']

            adf = row.adf_value
            if adf < 0.2:
                # start trading
                status = 'trading'
                if self.mid_ratio is None:
                    self.mid_ratio = row.close
                    self.step_gap = row.smooth_vol
                self.trading_logic(merged_data)

            if adf > 0.5:
                status = 'exiting'
                self.exit_logic(merged_data)
                self.mid_ratio = None



    def trading_logic(self,merged_mktdata):
        for day, row in merged_mktdata.iterrows():
            print(day)
            print(row)


        print(123)

    def waiting_logic(self,merged_mktdata):
        pass

    def exit_logic(self,merged_mktdata):
        pass















    def init(self):
        # setup algo paramters
        pass

    def on_bar_min(self):
        pass



if __name__ == '__main__':


    begin = datetime(2023, 5, 1)
    end = datetime(2024, 5, 30)

    exchange='binance'
    data_path = Path.cwd().parent / "data" / exchange / "full_data"

    symbol_1 = 'FILUSDT'
    symbol_2 = 'XLMUSDT'
    file_name = f"{symbol_1.upper()}-spot"
    file_path = data_path / f'{file_name}.csv'
    data = Mktdata(file_path, "1h", begin, end)
    symbol_object_1 = Symbol(symbol_1,data)

    file_name = f"{symbol_2.upper()}-spot"
    file_path = data_path / f'{file_name}.csv'
    data = Mktdata(file_path, "1h", begin, end)
    symbol_object_2 = Symbol(symbol_2,data)

    params = {
        "symbol_1":symbol_object_1,
        "symbol_2":symbol_object_2,
        "init_equity":10000,
        "amount": 1000,
        "steps":5,
        "gap":0.02,
        'init_ratio':1,
        'vol_window':20,
        'mean_window':10,
        'model_window':20,
        'trade_window':10,
    }

    algo = PairTrading('pair_trading', params, begin, end)
    algo.run()
    algo.display(True)


